package com.example.demo.algorithm.tree;

/**
 * @Author: xiatian
 * @CreateTime: 2022/2/3 1:31 下午
 * @Description:红黑树的实现
 */
public class RedBlackTree<Key extends Comparable<Key>, Value> {
    //根节点
    private Node root;
    //记录树中的元素的个数
    private int N;
    //红色链接
    private static final boolean RED = true;
    //黑色链接
    private static final boolean BLACK = false;

    public class Node {
        public Key key;
        private Value value;
        public Node left;
        public Node right;
        //由其父接待指向它的链接的颜色
        public boolean color;

        public Node(Key key, Value value, Node left, Node right, boolean color) {
            this.key = key;
            this.value = value;
            this.left = left;
            this.right = right;
        }
    }

    /**
     * 返回当前数的节点个数
     *
     * @return
     */
    private int size() {
        return N;
    }

    /**
     * 判断指向当前节点的链接的颜色
     *
     * @param x
     * @return
     */
    private boolean isRed(Node x) {
        if (x == null) {
            return false;
        }
        return x.color == RED;
    }

    /**
     * 左旋转
     *
     * @param h
     * @return
     */
    private Node rotateLeft(Node h) {
        //获取h节点的右子节点，表示为x
        Node x = h.right;
        //让x节点的左子节点成为h节点的右子节点
        h.right = x.left;
        //让h成为x节点的左子节点
        x.left = h;
        //让x节点的color属性等于h节点的color属性
        x.color = h.color;
        //让h节点的color属性变为红色
        h.color = RED;
        return x;
    }

    /**
     * 右旋
     *
     * @param h
     * @return
     */
    private Node rotateRight(Node h) {
        //获取h节点左子节点作为x
        Node x = h.left;
        //让x的右子节点作为h的左子节点
        h.left = x.right;
        //让h成为x的右子节点
        x.right = h;
        //让h的color成为x的color
        x.color = h.color;
        //让h的color成为红色
        h.color = RED;
        return x;
    }

    /**
     * 颜色反转，相当于完成拆分4-节点
     *
     * @param h
     */
    private void flipColors(Node h) {
        //当前节点变为红色
        h.color = RED;
        //左右节点变成黑色
        h.left.color = BLACK;
        h.right.color = BLACK;
    }

    /**
     * 在整个树上完成插入操作
     *
     * @param key
     * @param value
     */
    public void put(Key key, Value value) {
        root = put(root, key, value);
        //根节点的颜色总是黑色的
        root.color = BLACK;
    }

    /**
     * 在指定树中完成插入操作，并返回添加元素后的新树
     *
     * @param h
     * @param key
     * @param value
     * @return
     */
    public Node put(Node h, Key key, Value value) {
        //判断h是否为空，如果为空则直接返回一个红色的节点就可以了
        if (h == null) {
            //数量+1
            N++;
            return new Node(key, value, null, null, RED);
        }
        //比较h节点的键和key的大小
        int cmp = key.compareTo(h.key);
        if (cmp < 0) {
            //继续向左
            h.left = put(h.left, key, value);
        } else if (cmp > 0) {
            //继续向右
            h.right = put(h.right, key, value);
        } else {
            //值相等的替换
            h.value = value;
        }
        //进行左旋：当前节点h的左子节点是黑色，右子节点为红色，需要左旋
        if (isRed(h.right) && !isRed(h.left)) {
            h = rotateLeft(h);
        }
        //进行右旋(当前节点h的左子节点和左子节点的左子节点都为红色，需要右旋)
        if (isRed(h.left) && isRed(h.left.left)) {
            h = rotateRight(h);
        }
        //颜色反转:当前节点的左子节点和右子节点都为红色时，需要颜色反转
        if (isRed(h.right) && isRed(h.left)) {
            flipColors(h);
        }
        return h;
    }

    /**
     * 根据key从树中找出对相应的值
     *
     * @return
     */
    public Value get(Key key) {
        return get(root, key);
    }

    /**
     * 从指定的树x中，查找key对应的值
     *
     * @param x
     * @param key
     * @return
     */
    public Value get(Node x, Key key) {
        if (x == null) {
            return null;
        }
        //比较x节点的键和key的大小
        int cmp = key.compareTo(x.key);
        if (cmp > 0) {
            //如果key大于x节点的键，则继续找x节点的右子树
            return get(x.right, key);
        } else if (cmp < 0) {
            //如果key小于x节点的键，则继续找x节点的左子树
            return get(x.left, key);
        } else {
            //如果key等于x节点的键则返回值
            return x.value;
        }
    }

}
